from torch.utils.tensorboard.writer import SummaryWriter
import torchvision
from torch.utils.data import DataLoader

dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, transform=dataset_transform, download=True)
# 4个为一包，随机进行排列
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

# 测试数据集中的第一个样本
img, target = test_loader.dataset[0]
print(img.shape)
print(target)

writer = SummaryWriter("dataloader")
for epoch in range(2):
    global_step = 0
    for data in test_loader:
        imgs, targets = data
        writer.add_images("epoch: {}".format(epoch), imgs, global_step)
        global_step += 1
    
writer.close()